'''
Author: WangNing
Date: 2025-04-26 14:07:56
LastEditors: WangNing
LastEditTime: 2025-04-26 14:09:35
FilePath: /linetest/backend/utils/db.py
Description: 
Copyright (c) 2025 by VGE, All Rights Reserved. 
'''
'''
数据库连接模块
'''
import psycopg2
from psycopg2.extras import RealDictCursor
import time
from flask import g
from config.config import DB_CONFIG

def get_db(app):
    """获取数据库连接"""
    if 'db' not in g:
        retry_count = 0
        max_retries = 3
        while retry_count < max_retries:
            try:
                g.db = psycopg2.connect(**DB_CONFIG)
                # 设置自动提交为False，以便进行事务管理
                g.db.autocommit = False
                break
            except psycopg2.OperationalError as e:
                retry_count += 1
                if retry_count >= max_retries:
                    app.logger.error(f"数据库连接失败(尝试{retry_count}次): {str(e)}")
                    raise
                app.logger.warning(f"数据库连接失败，正在重试({retry_count}/{max_retries}): {str(e)}")
                time.sleep(1)  # 等待1秒后重试
    return g.db

def close_db(e=None):
    """关闭数据库连接"""
    db = g.pop('db', None)
    if db is not None:
        db.close()

def ensure_tables_exist(app):
    """确保数据库表已创建"""
    conn = None
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()
        
        # 检查projects表是否存在
        cursor.execute("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_name = 'projects'
            );
        """)
        if not cursor.fetchone()[0]:
            app.logger.warning("数据库表不存在，开始创建...")
            
            # 导入并执行数据库初始化脚本
            from init_database import create_tables
            create_tables()
            
            app.logger.info("数据库表创建完成")
        else:
            app.logger.info("数据库表已存在")
            
    except Exception as e:
        app.logger.error(f"检查数据库表时出错: {str(e)}")
    finally:
        if conn:
            conn.close() 